/* Poisson equation solver using red-black Gauss-Seidel successive over-relaxation iteration. */

# include "scene.h"
# include "discr.h"
# include <stdio.h>
# include <stdlib.h>
# include <math.h>

#ifdef _OMP_SOLVE_POISSON_
# include <omp.h>
#endif

/*
# include <time.h>
# include <sys/timeb.h>
*/

/* This structure is used to detect most sensitive cell in an itteration. */
struct dpmax
    {
     short im, jm, km, isecm;
     double dpm;
    };

extern void press_bc(struct grid *g, short tflag, unsigned short iRK);

extern long int solve_poisson(struct grid *g, double ****rhs, unsigned short iRK)
  {
   int cnt;
   short i, j, k, stop_cnt, tid, nTRD, isec, indx;
   double ***p, *r, *th, *z, ***d, *dp, coeff_p_ijk, _div2p, q;
   struct dpmax *tdp;
   struct cmtrx *cm;

// printf("in Solve_poisson()\n");
   th = g->th; stop_cnt = 0;

   nTRD = 1;
#  ifdef _OMP_SOLVE_POISSON_
   nTRD = g->nTRD;
#  endif

   tdp = (struct dpmax *)malloc(sizeof(struct dpmax)*nTRD);
   dp = (double *)malloc(sizeof(double)*nTRD);

   for(cnt = 1; cnt <= POISS_CNT_MAX; cnt++)
      {
       for(tid = 0; tid < nTRD; tid++)
          {
           tdp[tid].dpm = (double)0;
           tdp[tid].im = -1;
           tdp[tid].jm = -1;
           tdp[tid].km = -1;
           tdp[tid].isecm = -1;
          }

       for(isec = 0; isec < Nsec; isec++)
        {
         d = rhs[isec]; p = g->p[isec]; r = g->r[isec]; z = g->z[isec];
         tid = 0;

#        ifdef _OMP_SOLVE_POISSON_
         #pragma omp parallel private(i,j,k,tid,indx,q,coeff_p_ijk,_div2p)
#        endif
          {
        /* Initializing elements of the structure dpmax. */
#          ifdef _OMP_SOLVE_POISSON_
           tid = omp_get_thread_num();
#          endif

#          ifdef _OMP_SOLVE_POISSON_
           #pragma omp for schedule(dynamic,1)
#          endif
           for(k = NzG; k < NzG+g->nz[isec]; k++) for(i = NrG; i < NrG+g->nr[isec]; i++) for(j = NthG+(k+i)%2; j < NthG+g->nth; j += 2)
              {
               coeff_p_ijk = 0;
               _div2p = 0;

            /* Evaluating (1/r)*(dp/dr) at the cell (i,j,k) at cell-centers. */
               cm = &g->dcmx_rcc1[isec];
               for(q = 0.0, indx = 0; indx <= cm->iend[i-cm->respadi]-cm->istart[i-cm->respadi]; indx++)
                  if(cm->srcpadi+cm->istart[i-cm->respadi]+indx == i) coeff_p_ijk += 2*cm->mtrx[i-cm->respadi][indx]/(r[i]+r[i+1]);
                  else q += cm->mtrx[i-cm->respadi][indx]*p[cm->srcpadi+cm->istart[i-cm->respadi]+indx][j][k];

               _div2p += 2*q/(r[i]+r[i+1]);

            /* Evaluating d2p/dr2 at the cell (i,j,k) at cell-centers. */
               cm = &g->dcmx_rcc2[isec];
               for(q = 0.0, indx = 0; indx <= cm->iend[i-cm->respadi]-cm->istart[i-cm->respadi]; indx++)
                  if(cm->srcpadi+cm->istart[i-cm->respadi]+indx == i) coeff_p_ijk += cm->mtrx[i-cm->respadi][indx];
                  else q += cm->mtrx[i-cm->respadi][indx]*p[cm->srcpadi+cm->istart[i-cm->respadi]+indx][j][k];

               _div2p += q;

            /* Evaluating (1/r^2)*(d2p/dth2) at the cell (i,j,k). */
               cm = &g->dcmx_thcc2;
               for(q = 0.0, indx = 0; indx <= cm->iend[j-cm->respadi]-cm->istart[j-cm->respadi]; indx++)
                  if(cm->srcpadi+cm->istart[j-cm->respadi]+indx == j) coeff_p_ijk += 4*(cm->mtrx[j-cm->respadi][indx]/(r[i]+r[i+1]))/(r[i]+r[i+1]);
                  else q += cm->mtrx[j-cm->respadi][indx]*p[i][cm->srcpadi+cm->istart[j-cm->respadi]+indx][k];

               _div2p += 4*(q/(r[i]+r[i+1]))/(r[i]+r[i+1]);


            /* Evaluating d2p/dz2 at the cell (i,j,k). */
               cm = &g->dcmx_zcc2[isec][0];
               if(isec == 1) if( (i >= NrG+g->nr[0]) && (i < NrG+g->nr[1]-g->nr[2])) cm = &g->dcmx_zcc2[isec][1];
               for(q = 0.0, indx = 0; indx <= cm->iend[k-cm->respadi]-cm->istart[k-cm->respadi]; indx++)
                  if(cm->srcpadi+cm->istart[k-cm->respadi]+indx == k) coeff_p_ijk += cm->mtrx[k-cm->respadi][indx];
                  else q += cm->mtrx[k-cm->respadi][indx]*p[i][j][cm->srcpadi+cm->istart[k-cm->respadi]+indx];

               _div2p += q;

               dp[tid] = (d[i][j][k] - _div2p)/coeff_p_ijk - p[i][j][k];

//                dp[tid] = P_ijk_NEW - p[i][j][k];

// printf("%0.15E\t%0.15E\n",dp[tid],P_ijk_NEW - p[i][j][k]);

               if((dp[tid]<(double)0 ? -dp[tid]:dp[tid]) > (tdp[tid].dpm<(double)0 ? -tdp[tid].dpm:tdp[tid].dpm))
                  { tdp[tid].dpm = dp[tid]; tdp[tid].im = i; tdp[tid].jm = j; tdp[tid].km = k; tdp[tid].isecm = isec; }

               p[i][j][k] += PPE_RFACT*dp[tid];
              }

#          ifdef _OMP_SOLVE_POISSON_
           #pragma omp for schedule(dynamic,1)
#          endif
           for(k = NzG; k < NzG+g->nz[isec]; k++) for(i = NrG; i < NrG+g->nr[isec]; i++) for(j = NthG+(k+i+1)%2; j < NthG+g->nth; j += 2)
              {
               coeff_p_ijk = 0;
               _div2p = 0;

            /* Evaluating (1/r)*(dp/dr) at the cell (i,j,k) at cell-centers. */
               cm = &g->dcmx_rcc1[isec];
               for(q = 0.0, indx = 0; indx <= cm->iend[i-cm->respadi]-cm->istart[i-cm->respadi]; indx++)
                  if(cm->srcpadi+cm->istart[i-cm->respadi]+indx == i) coeff_p_ijk += 2*cm->mtrx[i-cm->respadi][indx]/(r[i]+r[i+1]);
                  else q += cm->mtrx[i-cm->respadi][indx]*p[cm->srcpadi+cm->istart[i-cm->respadi]+indx][j][k];

               _div2p += 2*q/(r[i]+r[i+1]);

            /* Evaluating d2p/dr2 at the cell (i,j,k) at cell-centers. */
               cm = &g->dcmx_rcc2[isec];
               for(q = 0.0, indx = 0; indx <= cm->iend[i-cm->respadi]-cm->istart[i-cm->respadi]; indx++)
                  if(cm->srcpadi+cm->istart[i-cm->respadi]+indx == i) coeff_p_ijk += cm->mtrx[i-cm->respadi][indx];
                  else q += cm->mtrx[i-cm->respadi][indx]*p[cm->srcpadi+cm->istart[i-cm->respadi]+indx][j][k];

               _div2p += q;

            /* Evaluating (1/r^2)*(d2p/dth2) at the cell (i,j,k). */
               cm = &g->dcmx_thcc2;
               for(q = 0.0, indx = 0; indx <= cm->iend[j-cm->respadi]-cm->istart[j-cm->respadi]; indx++)
                  if(cm->srcpadi+cm->istart[j-cm->respadi]+indx == j) coeff_p_ijk += 4*(cm->mtrx[j-cm->respadi][indx]/(r[i]+r[i+1]))/(r[i]+r[i+1]);
                  else q += cm->mtrx[j-cm->respadi][indx]*p[i][cm->srcpadi+cm->istart[j-cm->respadi]+indx][k];

               _div2p += 4*(q/(r[i]+r[i+1]))/(r[i]+r[i+1]);


            /* Evaluating d2p/dz2 at the cell (i,j,k). */
               cm = &g->dcmx_zcc2[isec][0];
               if(isec == 1) if( (i >= NrG+g->nr[0]) && (i < NrG+g->nr[1]-g->nr[2])) cm = &g->dcmx_zcc2[isec][1];
               for(q = 0.0, indx = 0; indx <= cm->iend[k-cm->respadi]-cm->istart[k-cm->respadi]; indx++)
                  if(cm->srcpadi+cm->istart[k-cm->respadi]+indx == k) coeff_p_ijk += cm->mtrx[k-cm->respadi][indx];
                  else q += cm->mtrx[k-cm->respadi][indx]*p[i][j][cm->srcpadi+cm->istart[k-cm->respadi]+indx];

               _div2p += q;

               dp[tid] = (d[i][j][k] - _div2p)/coeff_p_ijk - p[i][j][k];

//                dp[tid] = P_ijk_NEW - p[i][j][k];

               if((dp[tid]<(double)0 ? -dp[tid]:dp[tid]) > (tdp[tid].dpm<(double)0 ? -tdp[tid].dpm:tdp[tid].dpm))
                  { tdp[tid].dpm = dp[tid]; tdp[tid].im = i; tdp[tid].jm = j; tdp[tid].km = k; tdp[tid].isecm = isec; }

               p[i][j][k] += PPE_RFACT*dp[tid];
              }
          }
        }

       press_bc(g, 0, iRK);

    /* Locating the thread containing the maximum residual. */
       for(i = 0, tid = 1; tid < nTRD; tid++)
         if( (tdp[tid].dpm < (double)0 ? -tdp[tid].dpm:tdp[tid].dpm) > (tdp[i].dpm < (double)0 ? -tdp[i].dpm:tdp[i].dpm) ) i = tid;

       if(g->tptr != (FILE *)0) if( (cnt == 1) || (cnt%PPE_RESUAL_FREQ == 0) )
       fprintf(g->tptr, "=> %d: dpm = %E(%d, %d, %d | %d)\n", cnt, tdp[i].dpm, tdp[i].im, tdp[i].jm, tdp[i].km, tdp[i].isecm);

    /* Stop the iteration if the maximum residual is less than PRESS_DELTA continuously for POISS_CNT_MIN number of times. */
       if( (stop_cnt = ( (tdp[i].dpm < 0 ? -tdp[i].dpm:tdp[i].dpm) <= PRESS_DELTA) ? stop_cnt+1:0) == POISS_CNT_MIN) break;
      }

   if(g->tptr != (FILE *)0) fprintf(g->tptr, "N_ittr = %d; dpm = %E(%d, %d, %d | %d); %c\n",\
                    cnt, tdp[i].dpm, tdp[i].im, tdp[i].jm, tdp[i].km, tdp[i].isecm, (stop_cnt == POISS_CNT_MIN) ? 'Y':'N');
   free(dp); free(tdp);

   return cnt;
  }
